
# json->coco

import json
import os
import argparse
import shutil
# from collections import defaultdict
import cv2
import requests
from tqdm import tqdm
from PIL import Image  # 用于动态获取图像宽度和高度


def xuanze_predict(img_path,post_url):
    img = cv2.imread(img_path)
    # img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # img_x,img_y,_=img.shape

    success, encoded_image = cv2.imencode('.png', img)
    files = {
        'file': ('image.png', encoded_image.tobytes(), 'image/png')
    }

    response = requests.post(post_url, files=files, timeout=15)
    json_data=response.json()
    return json_data


def convert(size, box):
    """
    将 VOC 格式的边界框坐标转换为 YOLO 格式
    size: (width, height) 图像尺寸
    box: (xmin, xmax, ymin, ymax)
    """
    dw = 1. / size[0]
    dh = 1. / size[1]
    x_center = (box[0] + box[1]) / 2.0
    y_center = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x_center = x_center * dw
    w = w * dw
    y_center = y_center * dh
    h = h * dh
    return (x_center, y_center, w, h)

def save_yolo_label(label_path, labels):
    with open(label_path, 'w') as f:
        for cls_id, bbox in labels:
            line = f"{cls_id} {' '.join([f'{x:.6f}' for x in bbox])}\n"
            f.write(line)

def json_to_coco(input_dir, output_dir,post_url):
    """
    把接口结果转化成coco数据集格式
    :param input_dir: 图像文件的目录
    :param output_dir: 输出目录，用于保存yolo格式
    """
    classes=[]
    #创建类别名称文件
    images_path = os.path.join(output_dir, "images")
    if not os.path.exists(images_path):
        os.makedirs(images_path)
    labels_path = os.path.join(output_dir, "labels")
    if not os.path.exists(labels_path):
        os.makedirs(labels_path)


    # 遍历YOLO标注文件
    image_files = [f for f in os.listdir(input_dir) ]
    # print(annotation_files)
    for  image_file in tqdm(  image_files):
        if ".png" not in  image_file :
            continue
        image_path = os.path.join(input_dir, image_file)
        image_path_yolo = os.path.join(images_path, image_file)
        labels_path_yolo = os.path.join(labels_path, image_file.replace(".png",".txt"))

        # 检查图像文件是否存在
        if not os.path.exists(image_path):
            print(f"图像文件 {image_path} 不存在，跳过...")
            continue

        # 动态获取图像宽度和高度
        with Image.open(image_path) as img:
            image_width, image_height = img.size

        labels=[]
        # 读取检测接口结果
        try:
            image_res=xuanze_predict(image_path,post_url)
            for box in image_res["res"]["boxes"]:
                cls_id = box["cls_id"]  #yolo标签从0开始
                coordinate=box["coordinate"]
                bb=convert((image_width, image_height), [coordinate[0],coordinate[2],coordinate[1],coordinate[3]])
                labels.append((cls_id, bb))
            save_yolo_label(labels_path_yolo, labels)
            shutil.copy(image_path, image_path_yolo)
        except:
            print(f" 失败 {image_path}")


    print(f"转换完成！yolo格式的标注文件已保存到 {input_dir}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="把接口结果转化成coco数据集格式")
    parser.add_argument("--input_dir", type=str, required=False, help="图像文件的目录",default="/paddle/PaddleX/output/xuanze_test_img")
    parser.add_argument("--output_dir", type=str, required=False, help="输出目录，用于保存yolo格式文件",default="/paddle/PaddleX/dataset/yolo_data/yoyo_data")
    parser.add_argument("--post_url", type=str, required=False, help="调用接口",default="http://8.152.158.92:8080/predict")
    args = parser.parse_args([])

    json_to_coco(args.input_dir, args.output_dir,args.post_url)


